package com.pt25.base.util;

import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import okhttp3.*;
import okio.ByteString;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.PooledObjectFactory;
import org.apache.commons.pool2.impl.DefaultPooledObject;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

public class WsClientUtil2 {
    private static Map<String, GenericObjectPool<WebSocketClient>> webSocketClientPoolMap = new HashMap<>();

    public static GenericObjectPool<WebSocketClient> getWebsocketClientPool(WSConfig wsConfig) {
        String mapKey = wsConfig.getWebSocketUrl().split("\\?")[0];
        if (null == webSocketClientPoolMap.get(mapKey)) {
            GenericObjectPoolConfig<WebSocketClient> config = new GenericObjectPoolConfig<>();
            config.setMaxTotal(10);
            config.setMaxIdle(5);
            config.setMinIdle(1);
            config.setTestOnBorrow(true);
            config.setTestOnReturn(true);
            config.setTimeBetweenEvictionRunsMillis(30000); // 30 seconds
            config.setMinEvictableIdleTimeMillis(60000); // 60 seconds
            WebSocketClientFactory factory = new WebSocketClientFactory(wsConfig.getWebSocketUrl(), wsConfig.getHeaders());
            GenericObjectPool<WebSocketClient> pool = new GenericObjectPool<>(factory, config);
            webSocketClientPoolMap.put(mapKey, pool);
        }
        return webSocketClientPoolMap.get(mapKey);
    }

    @NoArgsConstructor
    @AllArgsConstructor
    @Builder
    @Data
    public static class WSConfig {
        private String webSocketUrl;
        @Builder.Default
        private Map<String, String> headers = new HashMap<>();
        @Builder.Default
        private Integer maxTotal = 20;
        @Builder.Default
        private Integer minIdle = 1;
        @Builder.Default
        private Integer maxIdle = 5;
        @Builder.Default
        private Boolean testOnBorrow = true;
        @Builder.Default
        private Boolean testOnReturn = true;
        @Builder.Default
        private Long timeBetweenEvictionRunsMillis = 30000L;//30秒
        @Builder.Default
        private Long minEvictableIdleTimeMillis = 60000L;//60秒
    }


    public static class WebSocketClient {
        private final OkHttpClient client;
        private WebSocket webSocket;
        private CountDownLatch latch;
        private MessageHandler messageHandler;

        public WebSocketClient() {
            client = new OkHttpClient();
        }

        public void connect(String url, Map<String, String> headers) throws InterruptedException {
            Request.Builder requestBuilder = new Request.Builder().url(url);
            headers.forEach(requestBuilder::addHeader);
            Request request = requestBuilder.build();
            latch = new CountDownLatch(1);
            webSocket = client.newWebSocket(request, new WebSocketListener() {
                @Override
                public void onOpen(WebSocket webSocket, Response response) {
                    latch.countDown();
                    if (messageHandler != null) {
                        messageHandler.handleOpen(webSocket, response);
                    }
                }

                @Override
                public void onMessage(WebSocket webSocket, String message) {
                    if (messageHandler != null) {
                        messageHandler.handleMessage(webSocket, message);
                    }
                }

                @Override
                public void onMessage(WebSocket webSocket, ByteString message) {
                    if (messageHandler != null) {
                        messageHandler.handleMessage(webSocket, message);
                    }
                }

                @Override
                public void onClosing(WebSocket webSocket, int code, String reason) {
                    webSocket.close(1000, null);
                    if (messageHandler != null) {
                        messageHandler.handleClosing(webSocket, code, reason);
                    }
                }

                @Override
                public void onClosed(WebSocket webSocket, int code, String reason) {
                    webSocket.close(1000, null);
                    if (messageHandler != null) {
                        messageHandler.handleClosed(webSocket, code, reason);
                    }
                }

                @Override
                public void onFailure(WebSocket webSocket, Throwable t, Response response) {
                    t.printStackTrace();
                    if (messageHandler != null) {
                        messageHandler.handleFailure(webSocket, t, response);
                    }
                }
            });
            latch.await(3, TimeUnit.SECONDS); // wait for connection to be established
        }

        public void sendMessage(String message) {
            webSocket.send(message);
        }


        public void sendMessage(ByteString message) {
            webSocket.send(message);
        }

        public void close() {
            if (webSocket != null) {
                webSocket.close(1000, "Goodbye");
            }
        }

        public void setMessageHandler(MessageHandler handler) {
            this.messageHandler = handler;
        }
    }

    public static class WebSocketClientFactory implements PooledObjectFactory<WebSocketClient> {
        private final String url;
        private final Map<String, String> headers;

        public WebSocketClientFactory(String url, Map<String, String> headers) {
            this.url = url;
            this.headers = headers;
        }

        @Override
        public PooledObject<WebSocketClient> makeObject() throws Exception {
            WebSocketClient client = new WebSocketClient();
            client.connect(url, headers);
            return new DefaultPooledObject<>(client);
        }

        @Override
        public void destroyObject(PooledObject<WebSocketClient> pooledObject) throws Exception {
            WebSocketClient client = pooledObject.getObject();
            client.close();
        }

        @Override
        public boolean validateObject(PooledObject<WebSocketClient> pooledObject) {
            WebSocketClient client = pooledObject.getObject();
            return client != null; // Add specific validation if necessary
        }

        @Override
        public void activateObject(PooledObject<WebSocketClient> pooledObject) throws Exception {
            // No activation required
        }

        @Override
        public void passivateObject(PooledObject<WebSocketClient> pooledObject) throws Exception {
            // No passivation required
        }
    }

    public interface MessageHandler {
        void handleMessage(WebSocket webSocket, String message);

        void handleMessage(WebSocket webSocket, ByteString message);

        void handleOpen(WebSocket webSocket, Response response);

        void handleClosing(WebSocket webSocket, int code, String reason);

        void handleClosed(WebSocket webSocket, int code, String reason);

        void handleFailure(WebSocket webSocket, Throwable t, Response response);
    }
}
